#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <curl/curl.h>
#include "cJSON.h"
#include "sds.h"
#include "cai.h"

struct cai_convo
{
    cJSON* messages;
    int max;
};

struct cai_convo* cai_convo_create(const char* prompt, int max)
{
    struct cai_convo* convo = malloc(sizeof(struct cai_convo));
    if (!convo)
    {
        return NULL;
    }

    convo->messages = cJSON_CreateArray();
    convo->max = max;

    if (prompt)
    {
        cJSON* prompt_msg = cJSON_CreateObject();
        cJSON_AddStringToObject(prompt_msg, "role", "system");
        cJSON_AddStringToObject(prompt_msg, "content", prompt);
        cJSON_AddItemToArray(convo->messages, prompt_msg);
    }

    return convo;
}

void cai_convo_destroy(struct cai_convo* convo)
{
    if (convo)
    {
        if (convo->messages)
        {
            cJSON_Delete(convo->messages);
        }
        free(convo);
    }
}

/**
 * @brief 加载文件内容到内存
 * 
 * @param path 文件路径
 * @param buffer 指向文件内容的指针地址
 * @return long 文件内容大小，失败返回-1
 */
static long load_file(const char *path, char **buffer)
{
    FILE *f = fopen(path, "rb");
    if (!f) {
        *buffer = NULL;
        return -1;
    }

    if (fseek(f, 0, SEEK_END) != 0) {
        fclose(f);
        *buffer = NULL;
        return -1;
    }

    // 获取文件大小
    long size = ftell(f);
    // 重置文件指针到文件开头
    fseek(f, 0, SEEK_SET);

    if (size == 0) {
        // 文件为空，无需分配内存
        *buffer = NULL;
        fclose(f);
        return 0;
    }

    *buffer = malloc(size);
    if (!*buffer) {
        fclose(f);
        *buffer = NULL;
        return -1;
    }

    size_t read = fread(*buffer, 1, size, f);
    if (read != (size_t)size) {
        free(*buffer);
        *buffer = NULL;
        fclose(f);
        return -1;
    }

    fclose(f);
    return size;
}

struct cai_convo* cai_convo_load(const char* file)
{
    if (!file)
    {
        return NULL;
    }

    char* jsonstr = NULL;
    long size = load_file(file, &jsonstr);
    if (size < 0)
    {
        return NULL;
    }

    struct cai_convo* convo = cai_convo_create(NULL, 0);
    if (!convo)
    {
        return NULL;
    }

    cJSON* root = cJSON_ParseWithLength(jsonstr, size);
    free(jsonstr);
    if (!root)
    {
        fprintf(stderr, "JSON parse failed: [%s]\n", cJSON_GetErrorPtr());
        cai_convo_destroy(convo);
        return NULL;
    }

    // 获取max字段
    cJSON* max = cJSON_GetObjectItem(root, "max");
    if (cJSON_IsNumber(max)) {
        convo->max = max->valueint;
    }

    cJSON* messages = cJSON_GetObjectItem(root, "messages");
    if (!cJSON_IsArray(messages))
    {
        fprintf(stderr, "cai_convo_load: messages is not an array\n");
        cai_convo_destroy(convo);
        cJSON_Delete(root);
        return NULL;
    }

    cJSON_Delete(convo->messages);  // 删除原有的空数组
    convo->messages = cJSON_Duplicate(messages, true);

    cJSON_Delete(root);

    return convo;
}

//文件扩展名和MIME类型映射表
static const struct {
    const char* ext;
    const char* mime;
} mime_types[] = {
    {".jpg", "image/jpeg"},
    {".jpeg", "image/jpeg"},
    {".png", "image/png"},
    {".gif", "image/gif"},
    {".mp3", "audio/mpeg"},
    {".wav", "audio/wav"},
    {".mp4", "video/mp4"},
};

#define ARRAY_SIZE(a) (sizeof(a) / sizeof(a[0]))

//根据文件扩展名获取MIME类型
static const char* get_mime_type(const char* path)
{
    if (!path)
    {
        return NULL;
    }

    const char* ext = strrchr(path, '.');
    if (!ext)
    {
        return NULL;
    }

    for (size_t i = 0; i < ARRAY_SIZE(mime_types); i++)
    {
        if (strcasecmp(ext, mime_types[i].ext) == 0)
        {
            return mime_types[i].mime;
        }
    }

    return NULL;
}

static cJSON* create_image_content(const char* mime, sds base64)
{
    cJSON* image_url = cJSON_CreateObject();
    sds url = sdscatfmt(sdsempty(), "data:%s;base64,%S", mime, base64);
    cJSON_AddStringToObject(image_url, "url", url);
    sdsfree(url);
    return image_url;
}

static cJSON* create_audio_content(const char* mime, sds base64)
{
    cJSON* audio_url = cJSON_CreateObject();
    sds data = sdsnew("data:;base64,");
    data = sdscat(data, base64);
    cJSON_AddStringToObject(audio_url, "data", data);
    sdsfree(data);
    char* subtype = strrchr(mime, '/');
    cJSON_AddStringToObject(audio_url, "format", subtype);
    return audio_url;
}

static cJSON* create_video_content(const char* mime, sds base64)
{
    cJSON* video_url = cJSON_CreateObject();
    sds url = sdsnew("data:;base64,");
    url = sdscat(url, base64);
    cJSON_AddStringToObject(video_url, "url", url);
    sdsfree(url);
    return video_url;
}

static const char base64_table[] = 
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

static sds b64encode(const char *path) {
    FILE *fp = fopen(path, "rb");
    if (!fp) return NULL;

    sds s = sdsempty();
    unsigned char buffer[3] = {0};
    int byte_count = 0;  // 当前缓存的字节数（0-3）
    int c;

    // 逐字节读取文件
    while ((c = fgetc(fp)) != EOF) {
        buffer[byte_count++] = (unsigned char)c;

        // 每满3字节进行编码
        if (byte_count == 3) {
            // 将3字节转换为4个Base64字符
            uint32_t triple = (buffer[0] << 16) | (buffer[1] << 8) | buffer[2];
            char quad[4] = {
                base64_table[(triple >> 18) & 0x3F],
                base64_table[(triple >> 12) & 0x3F],
                base64_table[(triple >> 6) & 0x3F],
                base64_table[triple & 0x3F]
            };

            // 追加到SDS
            s = sdscatlen(s, quad, 4);
            if (!s) break;  // 内存分配失败
            
            byte_count = 0;  // 重置缓存
        }
    }

    // 处理文件读取错误
    if (ferror(fp)) {
        sdsfree(s);
        fclose(fp);
        return NULL;
    }

    // 处理剩余字节
    if (byte_count > 0) {
        uint32_t partial = buffer[0] << 16;
        if (byte_count == 2) partial |= buffer[1] << 8;

        char quad[4] = {
            base64_table[(partial >> 18) & 0x3F],
            base64_table[(partial >> 12) & 0x3F],
            (byte_count == 2) ? base64_table[(partial >> 6) & 0x3F] : '=',
            '='
        };

        s = sdscatlen(s, quad, 4);
    }

    fclose(fp);
    return s;
}

static cJSON* create_content(const char* content, const char* file)
{
    cJSON* text = cJSON_CreateString(content);

    if (!file)
    {
        return text;
    }

    const char* mime = get_mime_type(file);
    if (!mime)
    {
        return text;
    }

    //将文件内容编码为base64字符串
    sds base64 = b64encode(file);
    if (!base64)
    {
        return text;
    }

    cJSON* file_content = cJSON_CreateObject();
    if (strncmp(mime, "image/", 6) == 0)
    {
        cJSON_AddStringToObject(file_content, "type", "image_url");
        cJSON_AddItemToObject(file_content, "image_url", create_image_content(mime, base64));
    }

    if (strncmp(mime, "audio/", 6) == 0)
    {
        cJSON_AddStringToObject(file_content, "type", "input_audio");
        cJSON_AddItemToObject(file_content, "input_audio", create_audio_content(mime, base64));
    }

    if (strncmp(mime, "video/", 6) == 0)
    {
        cJSON_AddStringToObject(file_content, "type", "video_url");
        cJSON_AddItemToObject(file_content, "video_url", create_video_content(mime, base64));
    }

    sdsfree(base64);

    cJSON* text_content = cJSON_CreateObject();
    cJSON_AddStringToObject(text_content, "type", "text");
    cJSON_AddStringToObject(text_content, "text", content);

    cJSON* contents = cJSON_CreateArray();
    cJSON_AddItemToArray(contents, text_content);
    cJSON_AddItemToArray(contents, file_content);

    return contents;
}

/**
 * @brief 向对话中增加一条消息
 * 
 * @param convo 对话结构体地址
 * @param role 消息的角色，可以是"system"、"user"、"assistant"、"tool"
 * @param content 消息的内容
 * @param file 文件路径
 * @return int 成功返回0，失败返回-1
 */
int cai_convo_append(struct cai_convo* convo,
                          const char* role,
                          const char* content,
                          const char* file)
{
    // 检查对话结构体和消息数组是否正确
    if (!convo || !cJSON_IsArray(convo->messages))
    {
        return -1;
    }

    // 如果对话中消息数量超过最大限制，则删除最旧的消息
    if (convo->max && cJSON_GetArraySize(convo->messages) >= convo->max)
    {
        cJSON_DeleteItemFromArray(convo->messages, 0);
    }

    cJSON* msg = cJSON_CreateObject();
    cJSON_AddStringToObject(msg, "role", role);
    cJSON_AddItemToObject(msg, "content", create_content(content, file));
    cJSON_AddItemToArray(convo->messages, msg);
    return 0;
}

int cai_convo_save(struct cai_convo* convo, const char* file)
{
    if (!file || !convo)
    {
        return -1;
    }

    cJSON* root = cJSON_CreateObject();
    cJSON_AddNumberToObject(root, "max", convo->max);
    cJSON_AddItemToObject(root, "messages", cJSON_Duplicate(convo->messages, true));
    char* jsonstr = cJSON_PrintUnformatted(root);
    if (!jsonstr)
    {
        return -1;
    }

    cJSON_Delete(root);

    FILE* f = fopen(file, "w");
    if (!f)
    {
        free(jsonstr);
        return -1;
    }

    if (fputs(jsonstr, f) == EOF)
    {
        fclose(f);
        free(jsonstr);
        return -1;
    }

    fclose(f);

    free(jsonstr);
    return 0;
}

struct cai_asst {
    char* api_key;
    char* base_url;
    char* model;
    char* resp_format; // 响应内容格式，默认为text
    float temperature; // 采样温度，控制模型生成文本的多样性，默认值为1.0
    cai_stream_callback stream_callback;
    void* stream_userdata;
    sds answer;
};

struct cai_asst* cai_asst_create(const char* api_key, const char* base_url, const char* model)
{
    if (!api_key || !base_url || !model) return NULL;

    struct cai_asst* asst = malloc(sizeof(struct cai_asst));
    if (!asst)
    {
        return NULL;
    }

    memset(asst, 0, sizeof(struct cai_asst));

    asst->api_key = strdup(api_key);
    asst->base_url= strdup(base_url);
    asst->model = strdup(model);
    asst->answer = sdsempty();
    asst->resp_format = strdup("text");
    asst->temperature = 1.0;
    return asst;
}

void cai_asst_set_temperature(struct cai_asst* asst, float temperature)
{
    asst->temperature = temperature;
}

void cai_asst_set_format(struct cai_asst* asst, const char* format)
{
    if (format)
    {
        asst->resp_format = strdup(format);
    }
}

void cai_asst_destroy(struct cai_asst* asst)
{
    if (asst)
    {
        if (asst->api_key)
        {
            free(asst->api_key);
        }
        if (asst->base_url)
        {
            free(asst->base_url);
        }
        if (asst->model)
        {
            free(asst->model);
        }
        if (asst->resp_format)
        {
            free(asst->resp_format);
        }
        sdsfree(asst->answer);
        free(asst);
    }
}

struct rbuf {
    sds response;
    int stream;
    struct cai_asst* asst;
};

/* SSE 事件处理函数 */
static void process_sse(struct cai_asst *asst, char *sse)
{
    if (strncmp(sse, "data:", 5) == 0) {
        char* jsonstr = sse + 5;
        if (strstr(jsonstr, "[DONE]")) {
            if (asst->stream_callback)
            {
                asst->stream_callback(asst->stream_userdata, NULL, NULL);
            }
            return;
        }

        cJSON* root = cJSON_Parse(jsonstr);
        if (!root) {
            fprintf(stderr, "JSON parse failed: [%s]\n", cJSON_GetErrorPtr());
            return;
        }

        cJSON* choices = cJSON_GetObjectItem(root, "choices");
        if (cJSON_IsArray(choices)) {
            cJSON* choice = cJSON_GetArrayItem(choices, 0);
            if (choice) {
                cJSON* delta = cJSON_GetObjectItem(choice, "delta");
                if (delta) {
                    cJSON* content = cJSON_GetObjectItem(delta, "content");
                    if (content && content->valuestring) {
                        asst->answer = sdscat(asst->answer, content->valuestring);
                        
                        if (asst->stream_callback)
                        {
                            asst->stream_callback(asst->stream_userdata, content->valuestring, NULL);
                        }
                    }

                    cJSON* reason = cJSON_GetObjectItem(delta, "reasoning_content");
                    if (cJSON_IsString(reason))
                    {
                        if (asst->stream_callback)
                        {
                            asst->stream_callback(asst->stream_userdata, NULL, reason->valuestring);
                        }
                    }
                }
            }
        }

        cJSON_Delete(root);
    }
}

static size_t on_header(char *data, size_t size, size_t nmemb, void *clientp) {
    //puts(data);
    struct rbuf *buf = (struct rbuf *)clientp;
    //如果是Content-Type字段，则检查字段的值是否包含"text/event-stream"
    if (strncasecmp(data, "Content-Type:", 13) == 0) {
        if (strstr(data, "text/event-stream") != NULL) {
            buf->stream = 1;
        }
    }
    return size * nmemb;
}

static size_t on_reply(char *data, size_t size, size_t nmemb, void *clientp) {
    size_t realsize = size * nmemb;
    struct rbuf *buf = (struct rbuf *)clientp;

    //将收到的数据保存到缓冲区
    buf->response = sdscatlen(buf->response, data, realsize);

    //如果不是流式数据，则直接返回，等响应消息接收完成后统一处理
    if (!buf->stream)
    {
        return realsize;
    }

    //如果SSE没有收完，则直接返回
    char* end;
    while(end = strstr(buf->response, "\n\n"))
    {
        sds sse = sdsempty();
        sse = sdscpylen(sse, buf->response, end - buf->response);
        process_sse(buf->asst, sse);
        sdsfree(sse);
    
        //将已处理的SSE数据从缓冲区中删除
        sdsrange(buf->response, (end - buf->response) + 2, -1);
    }

    return realsize;
}

//发送POST请求
//url: 请求地址
//headers: 增加的请求头部字段
//data: 请求正文
//arg:  透传参数
//返回值: 成功返回0，失败返回-1
static int http_post(char* url,
                     struct curl_slist* headers, 
                     char* data, 
                     void* arg)
{
    struct rbuf* buf = arg;

    CURL *client = curl_easy_init();
    if (!client)
    {
        return -1;
    }

    curl_easy_setopt(client, CURLOPT_URL, url);
    curl_easy_setopt(client, CURLOPT_POST, 1);
    curl_easy_setopt(client, CURLOPT_POSTFIELDS, data);

    if (headers)
    {
        curl_easy_setopt(client, CURLOPT_HTTPHEADER, headers);
    }

    // 将服务器返回的响应报文保存到文件流中
    curl_easy_setopt(client, CURLOPT_WRITEDATA, buf);
    curl_easy_setopt(client, CURLOPT_WRITEFUNCTION, on_reply);
    curl_easy_setopt(client, CURLOPT_HEADERDATA, buf);
    curl_easy_setopt(client, CURLOPT_HEADERFUNCTION, on_header);

    CURLcode error = curl_easy_perform(client);
    if (error != CURLE_OK)
    {
        fprintf(stderr, "curl_easy_perform: %s\n", curl_easy_strerror(error));
        curl_easy_cleanup(client);
        return -1;
    }

    long status = 0;
    curl_easy_getinfo(client, CURLINFO_RESPONSE_CODE, &status);
    if (status != 200)
    {
        fprintf(stderr, "HTTP status code: %ld\n", status);
    }

    curl_easy_cleanup(client);

    return 0;
}

/**
 * @brief 发送请求消息
 * 
 * @param asst AI助手结构体地址
 * @param history 对话历史记录
 * @return char* 非流式响应返回响应消息正文，流式响应或出错返回NULL
 */
static char* send_request(struct cai_asst* asst, cJSON* history)
{
    if (!asst->api_key || !asst->base_url)
    {
        return NULL;
    }

    struct rbuf buf;
    buf.response = sdsempty();
    buf.asst = asst;
    buf.stream = 0;

    sds auth = sdsnew("Authorization: Bearer ");
    auth = sdscat(auth, asst->api_key);

    sds url = sdsnew(asst->base_url);
    url = sdscat(url, "/chat/completions");

    // 设置 HTTP 请求头
    struct curl_slist* headers = NULL;
    headers = curl_slist_append(headers, "Content-Type: application/json");
    headers = curl_slist_append(headers, auth);

    // 构建请求体 JSON 结构
    cJSON* obj = cJSON_CreateObject();
    cJSON_AddStringToObject(obj, "model", asst->model);
    cJSON* history_copy = cJSON_Duplicate(history, true);
    cJSON_AddItemToObject(obj, "messages", history_copy);
    if (asst->stream_callback)
    {
        cJSON_AddBoolToObject(obj, "stream", true);
    }
    cJSON_AddNumberToObject(obj, "temperature", asst->temperature);

    cJSON* response_format = cJSON_CreateObject();
    cJSON_AddStringToObject(response_format, "type", asst->resp_format);
    cJSON_AddItemToObject(obj, "response_format", response_format);

    char* json = cJSON_Print(obj);  // 序列化 JSON 字符串
    cJSON_Delete(obj);

    //puts(json);

    // 发送 HTTP POST 请求并获取响应
    size_t size = strlen(json);
    if (http_post(url, headers, json, &buf) < 0)
    {
        free(json);
        curl_slist_free_all(headers);
        sdsfree(auth);
        sdsfree(url);
        sdsfree(buf.response);
        return NULL;
    }

    // 清理请求资源
    free(json);
    curl_slist_free_all(headers);
    sdsfree(auth);
    sdsfree(url);

    //流式响应返回NULL
    char* response = NULL;

    if (!buf.stream)
    {
        response = strdup(buf.response);
    }

    sdsfree(buf.response);

    return response;
}

static void process_response(struct cai_asst* asst, char* response)
{
    if (!response)
    {
        return;
    }

    cJSON* obj = cJSON_Parse(response);
    if (!obj)
    {
        fprintf(stderr, "JSON parse failed: [%s]\n", cJSON_GetErrorPtr());
        return;
    }

    // 提取有效响应内容
    cJSON* choices = cJSON_GetObjectItem(obj, "choices");

    /* 增强的错误处理逻辑：检查 API 返回的错误信息结构
        * 处理两种情况：
        * 1. 空响应数组（choices 为空）
        * 2. 显式的错误对象（error 字段）
        */
    if (!choices || cJSON_GetArraySize(choices) == 0) {
        cJSON* error = cJSON_GetObjectItem(obj, "error");
        if (error) {
            cJSON* error_msg = cJSON_GetObjectItem(error, "message");
            if (error_msg && error_msg->valuestring) {
                fprintf(stderr, "%s\n", error_msg->valuestring);
            }
        }
    }

    // 提取对话结果并更新历史记录
    cJSON* choice = cJSON_GetArrayItem(choices, 0);
    if (choice) {
        cJSON* message = cJSON_GetObjectItem(choice, "message");
        if (message) {
            cJSON* content = cJSON_GetObjectItem(message, "content");
            if (content && content->valuestring) {
                asst->answer = sdscat(asst->answer, content->valuestring);
            }
        }
    }

    cJSON_Delete(obj);
}

char* cai_asst_chat(struct cai_asst* asst, 
                    struct cai_convo* convo, 
                    const char* question,
                    const char* file)
{
    struct cai_convo* new_convo;
    if (!asst)
    {
        return NULL;
    }

    if (!convo)
    {
        new_convo = cai_convo_create(NULL, 0);
    }
    else
    {
        new_convo = convo;
    }

    cai_convo_append(new_convo, "user", question, file);

    //清空上一次的回答
    sdsfree(asst->answer);
    asst->answer = sdsempty();

    char* response = send_request(asst, new_convo->messages);
    if (response)
    {
        process_response(asst, response);
    }

    cai_convo_append(new_convo, "assistant", asst->answer, NULL);

    if (!convo)
    {
        cai_convo_destroy(new_convo);
    }

    return asst->answer;
}

void cai_set_stream_callback(struct cai_asst* asst, cai_stream_callback callback, void* userdata)
{
    asst->stream_callback = callback;
    asst->stream_userdata = userdata;
}